package com.demo.search.engine.annotation;

import java.io.IOException;
import java.io.StringReader;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.highlight.Highlighter;
import org.apache.lucene.search.highlight.InvalidTokenOffsetsException;
import org.apache.lucene.search.highlight.QueryScorer;
import org.apache.lucene.search.highlight.SimpleHTMLFormatter;
import org.apache.lucene.search.vectorhighlight.BaseFragmentsBuilder;
import org.apache.lucene.search.vectorhighlight.FastVectorHighlighter;
import org.apache.lucene.search.vectorhighlight.FieldQuery;
import org.apache.lucene.search.vectorhighlight.FragListBuilder;
import org.apache.lucene.search.vectorhighlight.FragmentsBuilder;
import org.apache.lucene.search.vectorhighlight.ScoreOrderFragmentsBuilder;
import org.apache.lucene.search.vectorhighlight.SimpleFragListBuilder;

import com.demo.common.StringUtils;
import com.demo.search.engine.HighlighterHandler;
import com.demo.search.model.FieldModel;
import com.demo.search.util.LuceneUtils;

/**
 * 高亮处理默认实现
 * 
 * @author yuu.kk
 */
public class HighlighterHandlerImpl implements HighlighterHandler {

	private FastVectorHighlighter highlighter;
	private Analyzer analyzer;
	private Query query;
	private IndexReader reader;
	private String[] preColorTags = new String[]{ "<font color=\"red\">" };
	private String[] postColorTags = new String[]{ "</font>" };
	/**
	 * @param query
	 * @param reader
	 * @param preColorTags
	 * @param postColorTags
	 */
	public HighlighterHandlerImpl(Analyzer analyzer, Query query, IndexReader reader, String[] preColorTags, String[] postColorTags) {
		this.analyzer = analyzer;
		this.query = query;
		this.reader = reader;
		//this.highlighter = getHightlighter(preColorTags, postColorTags);
		this.preColorTags = preColorTags;
		this.postColorTags = postColorTags;
	}
	
	/**
	 * 构造函数
	 * @param query
	 * @param reader
	 */
	public HighlighterHandlerImpl(Analyzer analyzer, Query query, IndexReader reader) {
		this.analyzer = analyzer;
		this.query = query;
		this.reader = reader;
		//this.highlighter = getHightlighter();
	}
	
	/**
	 * 构造函数
	 * @param query
	 * @param reader
	 * @param highlighter
	 */
	public HighlighterHandlerImpl(Analyzer analyzer, Query query, IndexReader reader, FastVectorHighlighter highlighter) {
		this.analyzer = analyzer;
		this.query = query;
		this.reader = reader;
		this.highlighter = highlighter;
	}
	
	@Override
	public boolean support(FieldModel model) {
		return model.allowHighlighter();
	}

	@Override
	public String render(Document doc, int docID, FieldModel model) {

		String highLightResult = null;
		FieldQuery fieldQuery = highlighter.getFieldQuery(query);
		// highlighter.setTextFragmenter(new SimpleFragmenter(100)); 
		if (highlighter == null) {
			highlighter = getHightlighter(preColorTags, postColorTags);
		}
		try {
			highLightResult = highlighter.getBestFragment(fieldQuery, reader, docID, model.getName(), 100);
		} catch (IOException e) {
			e.printStackTrace();
		}
		return highLightResult;
	}

	@Override
	public String render(String content, FieldModel model) {

		String preTag = preColorTags[0];
		String postTag = postColorTags[0];
		if (StringUtils.isEmpty(content)) {
			return "";
		}
		String highLightResult = null;
		// "<font color=\"red\">" - "</font>"
		SimpleHTMLFormatter simpleHtmlFormatter = new SimpleHTMLFormatter(preTag, postTag);
		Highlighter highlighter = new Highlighter(simpleHtmlFormatter, new QueryScorer(query));
		//QueryScorer
		try {
			if (analyzer == null) {
				analyzer = LuceneUtils.createChineseAnalyzer();
			}
			TokenStream tokenStream = analyzer.tokenStream(model.getName(), new StringReader(content));
			String highlighterStr = highlighter.getBestFragment(tokenStream, content);
			return highlighterStr == null ? content : highlighterStr;
		} catch (IOException e) {
			e.printStackTrace();
		} catch (InvalidTokenOffsetsException e) {
			e.printStackTrace();
		}
		return highLightResult;
	}
	
	/**
	 * 获得高亮标记对象
	 * @return FastVectorHighlighter
	 */
	private FastVectorHighlighter getHightlighter() {
		return getHightlighter(this.preColorTags, this.postColorTags);
	}
	
	/**
	 * 获得高亮标记对象
	 * @return FastVectorHighlighter
	 */
	private FastVectorHighlighter getHightlighter(String[] preColorTags, String[] postColorTags) {
		// new Highlighter(simpleHTMLFormatter, new QueryScorer(query));
		FragListBuilder fragListBuilder = new SimpleFragListBuilder();
		FragmentsBuilder fragmentBuilder = new ScoreOrderFragmentsBuilder(BaseFragmentsBuilder.COLORED_PRE_TAGS, BaseFragmentsBuilder.COLORED_POST_TAGS);
		
		return new FastVectorHighlighter(true, true, fragListBuilder, fragmentBuilder);
	}
}
